from typing import List, Dict, Generator, Any, Optional
import requests
import json
from pydantic import BaseModel
from fastapi import HTTPException
import logging
import re

# 初始化日志配置
logging.basicConfig(
    level=logging.DEBUG,  # 改为 INFO 级别，减少详细日志
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("RagflowPipeline")

class Pipeline:
    class Valves(BaseModel):
        ragflow_base_url: str = "http://host.docker.internal:8888"
        ragflow_api_key: str = "ragflow-UwM2I5MjBhOGUyMTExZjBhNTRjNTZkNj"
        default_model: str = "deepseek-r1:32b"
        chat_id: str = "d0604f508fa311f0beb816405564b0b0"
        session_id: str = "d3b13ec5300a4fdebac5b62ffbeb2682"
        auto_create_chat: bool = False
        stream_response: bool = True

    def __init__(self):
        self.valves = self.Valves()
        self.active_chat_id: Optional[str] = None
        self.headers: Dict[str, str] = {}
        self.current_session: str = ""
        logger.info("RagflowPipeline 初始化完成")
        
        if not hasattr(self.valves, 'chat_id'):
            logger.error("Valves类中未定义chat_id属性！")
            raise AttributeError("Valves类必须定义chat_id属性")
        if not hasattr(self.valves, 'session_id'):
            logger.error("Valves类中未定义session_id属性！")
            raise AttributeError("Valves类必须定义session_id属性")

    async def on_startup(self):
        logger.info(f"RagflowPipeline 启动")
        pass

    async def on_shutdown(self):
        logger.info(f"RagflowPipeline 关闭")
        pass
        
    def _init_headers(self):
        """初始化请求头"""
        masked_key = self.valves.ragflow_api_key[:10] + "..." if self.valves.ragflow_api_key else "None"
        logger.info(f"API Key: {masked_key}")
        
        if self.valves.ragflow_api_key:
            self.headers = {
                "Authorization": f"Bearer {self.valves.ragflow_api_key}",
                "Content-Type": "application/json"
            }

    def pipe(
        self,
        user_message: str,
        model_id: str,
        messages: List[Dict[str, str]],
        body: Dict[str, Any]
    ) -> Generator[str, None, None]:
        logger.info("===== 开始处理用户请求 =====")
        try:
            # 识别请求类型
            request_type = self._identify_request_type(body, user_message, messages)
            logger.info(f"请求类型: {request_type}")
            
            # 根据请求类型处理
            if request_type == "generate_tags":
                yield from self._handle_tag_generation(user_message, messages)
                return
            elif request_type == "generate_title":
                yield from self._handle_title_generation(user_message, messages)
                return
            elif request_type == "generate_follow_ups":
                yield from self._handle_follow_up_generation(user_message, messages)
                return
            else:
                yield from self._handle_chat_request(user_message, body)
                return

        except Exception as e:
            logger.error(f"处理请求失败: {str(e)}", exc_info=True)
            yield f"处理请求时发生错误: {str(e)}"
        finally:
            logger.info("===== 请求处理结束 =====")

    def _identify_request_type(self, body: Dict[str, Any], user_message: str, messages: List[Dict[str, str]]) -> str:
        """
        识别 Open-WebUI 的请求类型
        """
        # 检查 body 中是否有指示请求类型的字段
        if body.get("task") == "generate_tags":
            return "generate_tags"
        elif body.get("task") == "generate_title":
            return "generate_title"
        elif body.get("task") == "generate_follow_ups":
            return "generate_follow_ups"
        
        # 检查消息内容是否包含特定模式
        if "生成标签" in user_message or "tag" in user_message.lower():
            return "generate_tags"
        elif "生成标题" in user_message or "title" in user_message.lower():
            return "generate_title"
        elif "生成后续问题" in user_message or "follow up" in user_message.lower():
            return "generate_follow_ups"
        
        # 默认处理为聊天请求
        return "chat"

    def _handle_chat_request(self, user_message: str, body: Dict[str, Any]) -> Generator[str, None, None]:
        """处理普通聊天请求"""
        self._init_headers()
        logger.info(f"处理聊天请求: {user_message}")
        
        # 使用配置的聊天和会话ID
        self.active_chat_id = self.valves.chat_id
        self.current_session = self.valves.session_id
        logger.info(f"聊天ID: {self.active_chat_id}")
        logger.info(f"会话ID: {self.current_session}")

        if not self.active_chat_id:
            logger.error("未找到活跃对话")
            yield "错误: 未配置聊天ID"
            return

        # 修改提示，明确告诉模型不要输出思考过程
        enhanced_prompt = f"""
        {user_message}
        
        重要指令：
        1. 请直接回答问题，不要包含任何思考过程或<think>标签
        2. 不要解释你的思考过程，只需提供最终答案
        3. 如果可能，请用简洁明了的语言回答
        """
        
        # 发送对话请求
        payload = {
            "question": enhanced_prompt,
            "stream": self.valves.stream_response,
            "session_id": self.current_session
        }
        logger.debug(f"请求体: {json.dumps(payload, ensure_ascii=False)}")

        # 构建URL并清理可能的不可见字符
        url = f"{self.valves.ragflow_base_url}/api/v1/chats/{self.active_chat_id}/completions"
        url = ''.join(char for char in url if char.isprintable())
        logger.info(f"发送请求到: {url}")

        # 发送请求到RAGFlow
        timeout = 120
        try:
            response = requests.post(
                url,
                headers=self.headers,
                json=payload,
                stream=self.valves.stream_response,
                timeout=timeout
            )
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            logger.error(f"请求失败: {e}")
            if hasattr(e, 'response') and e.response is not None:
                logger.error(f"响应内容: {e.response.text}")
            yield f"请求失败: {str(e)}"
            return

        logger.info(f"响应状态码: {response.status_code}")
        logger.info(f"Content-Type: {response.headers.get('Content-Type', '未知')}")
        logger.info(f"Transfer-Encoding: {response.headers.get('Transfer-Encoding', '未知')}")

        # 处理响应
        yield from self._process_response(response, "chat")

    def _handle_tag_generation(self, user_message: str, messages: List[Dict[str, str]]) -> Generator[str, None, None]:
        """处理标签生成请求"""
        logger.info("处理标签生成请求")
        
        # 提取对话历史
        conversation = self._extract_conversation(messages)
        
        # 构建标签生成提示
        prompt = f"""
        请为以下对话生成1-3个广义标签以及更具体的子主题标签。
        对话内容：{conversation}
        
        请只返回JSON格式的结果，不要包含任何其他文本或思考过程。
        格式必须严格遵循：
        {{
            "tags": ["标签1", "标签2", "标签3"]
        }}
        """
        
        # 发送到RAGFlow
        yield from self._send_custom_prompt(prompt, "generate_tags")

    def _handle_title_generation(self, user_message: str, messages: List[Dict[str, str]]) -> Generator[str, None, None]:
        """处理标题生成请求"""
        logger.info("处理标题生成请求")
        
        # 提取对话历史
        conversation = self._extract_conversation(messages)
        
        # 构建标题生成提示
        prompt = f"""
        请为以下对话生成一个简洁的标题。
        对话内容：{conversation}
        
        请直接返回标题文本，不需要其他格式。
        """
        
        # 发送到RAGFlow
        yield from self._send_custom_prompt(prompt, "generate_title")

    def _handle_follow_up_generation(self, user_message: str, messages: List[Dict[str, str]]) -> Generator[str, None, None]:
        """处理后续问题生成请求"""
        logger.info("处理后续问题生成请求")
        
        # 提取对话历史
        conversation = self._extract_conversation(messages)
        
        # 构建后续问题生成提示
        prompt = f"""
        请为以下对话生成3个相关的后续问题。
        对话内容：{conversation}
        
        请只返回JSON格式的结果，不要包含任何其他文本或思考过程。
        格式必须严格遵循：
        {{
            "suggestions": ["问题1", "问题2", "问题3"]
        }}
        """
        
        # 发送到RAGFlow
        yield from self._send_custom_prompt(prompt, "generate_follow_ups")

    def _extract_conversation(self, messages: List[Dict[str, str]]) -> str:
        """从消息列表中提取对话内容"""
        conversation = []
        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role and content:
                conversation.append(f"{role}: {content}")
        return "\n".join(conversation)

    def _send_custom_prompt(self, prompt: str, request_type: str) -> Generator[str, None, None]:
        """发送自定义提示到RAGFlow"""
        self._init_headers()
        
        # 使用配置的聊天和会话ID
        self.active_chat_id = self.valves.chat_id
        self.current_session = self.valves.session_id
        
        # 发送自定义提示
        payload = {
            "question": prompt,
            "stream": self.valves.stream_response,
            "session_id": self.current_session
        }
        
        # 构建URL
        url = f"{self.valves.ragflow_base_url}/api/v1/chats/{self.active_chat_id}/completions"
        url = ''.join(char for char in url if char.isprintable())
        logger.info(f"发送请求到: {url}")
        
        # 发送请求到RAGFlow
        try:
            response = requests.post(
                url,
                headers=self.headers,
                json=payload,
                stream=self.valves.stream_response,
                timeout=300
            )
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            logger.error(f"请求失败: {e}")
            yield f"请求失败: {str(e)}"
            return
        
        logger.info(f"响应状态码: {response.status_code}")
        
        # 处理响应
        yield from self._process_response(response, request_type)

    def _process_streaming_answer(self, answer: str, in_think_block: bool, think_buffer: str, final_answer: str) -> Dict[str, Any]:
        """实时处理流式回答中的思考内容"""
        output = ""
        remaining_content = answer
        
        # 如果当前在思考块中，先处理思考块
        if in_think_block:
            if "</think>" in remaining_content:
                # 找到思考块结束
                think_end = remaining_content.find("</think>") + len("</think>")
                think_buffer += remaining_content[:think_end]
                remaining_content = remaining_content[think_end:]
                in_think_block = False
                
                # 只保留最后一次思考内容
                final_answer = think_buffer
                think_buffer = ""
            else:
                # 整个内容都在思考块中
                think_buffer += remaining_content
                remaining_content = ""
        
        # 处理剩余内容中的思考块
        while remaining_content:
            think_start = remaining_content.find("<think>")
            
            if think_start == -1:
                # 没有思考块，全部是非思考内容
                output += remaining_content
                break
            
            # 提取思考块之前的内容
            output += remaining_content[:think_start]
            remaining_content = remaining_content[think_start:]
            
            # 处理思考块
            think_end = remaining_content.find("</think>")
            
            if think_end == -1:
                # 思考块没有结束，标记为在思考块中
                think_buffer = remaining_content
                in_think_block = True
                remaining_content = ""
            else:
                # 思考块完整
                think_end += len("</think>")
                think_content = remaining_content[:think_end]
                remaining_content = remaining_content[think_end:]
                
                # 只保留最后一次思考内容
                final_answer = think_content
        
        return {
            "in_think_block": in_think_block,
            "think_buffer": think_buffer,
            "final_answer": final_answer,
            "output": output
        }

    def _process_response(self, response, request_type: str = "chat") -> Generator[str, None, None]:
        """处理RAGFlow响应"""
        # 处理流式响应
        if self.valves.stream_response and "text/event-stream" in response.headers.get('Content-Type', ''):
            logger.info("处理流式响应 (SSE)")
            buffer = ""
            in_think_block = False  # 标记是否在思考块中
            think_buffer = ""  # 用于累积思考内容
            final_answer = ""  # 用于累积最终答案
            
            for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
                if chunk:
                    logger.debug(f"收到块: {repr(chunk)}")
                    buffer += chunk
                    
                    # 处理缓冲区中的完整行
                    while '\n' in buffer:
                        line, buffer = buffer.split('\n', 1)
                        line = line.strip()
                        
                        if not line:
                            continue
                            
                        # 处理SSE格式
                        if line.startswith('data:'):
                            data_content = line[5:].strip()
                            
                            if data_content == '[DONE]':
                                logger.debug("收到流结束标记")
                                # 处理最终答案
                                if final_answer:
                                    # 清理最终答案中的思考内容
                                    cleaned_final = self._clean_ragflow_answer(final_answer, request_type)
                                    if cleaned_final:
                                        yield cleaned_final
                                continue
                                
                            if not data_content:
                                continue
                                
                            try:
                                data = json.loads(data_content)
                                logger.info(f"解析的数据: {data}")
                                is_final_answer = False
                                
                                # 提取答案
                                if isinstance(data, dict) and data.get("code") == 0:
                                    if isinstance(data, dict):
                                        answer_data = data.get("data", {})
                                        if isinstance(data, dict) and data.get("code") == 0:
                                            data_internal = data.get("data")
                                            if isinstance(data_internal, dict):
                                                reference_data = data_internal.get("reference", {})
                                                if isinstance(reference_data, dict):
                                                    total_data = reference_data.get("total")
                                                    logger.info(f"total_data: {total_data}")
                                                    if total_data is not None:
                                                        if total_data > 0:
                                                            is_final_answer = True
                                                            logger.info(f"is_final_answer: {is_final_answer}")
                                            
                                        if isinstance(answer_data, dict):
                                            answer = answer_data.get("answer", "")
                                        elif (type(answer_data) is bool):
                                            answer = answer_data
                                        else:
                                            answer = ""  # 或者处理非字典情况
                                            logger.warning(f"Unexpected data type for answer_data: {type(answer_data)}")
                                    else:
                                        answer = ""
                                        logger.warning(f"Unexpected data type for data: {type(data)}")
                                        
                                    if answer and is_final_answer is True:
                                        logger.info(f"The final answer: {answer}")
                                        if answer:
                                            yield answer
                                else:
                                    error_msg = data.get("message", "未知错误")
                                    logger.error(f"RAGFlow错误: {error_msg}")
                                    yield f"错误: {error_msg}"
                            except json.JSONDecodeError as e:
                                logger.warning(f"JSON解析错误: {e}, 内容: {data_content}")
                                # 如果不是JSON，尝试直接返回（但先过滤思考内容）
                                if data_content.strip():
                                    # 实时处理思考内容
                                    processed_answer = self._process_streaming_answer(
                                        data_content, in_think_block, think_buffer, final_answer
                                    )
                                    
                                    in_think_block = processed_answer["in_think_block"]
                                    think_buffer = processed_answer["think_buffer"]
                                    final_answer = processed_answer["final_answer"]
                                    
                                    # 如果有非思考内容，立即返回
                                    if processed_answer["output"]:
                                        yield processed_answer["output"]
            
            # 处理缓冲区中剩余的内容
            if buffer.strip():
                # 实时处理思考内容
                processed_answer = self._process_streaming_answer(
                    buffer, in_think_block, think_buffer, final_answer
                )
                
                in_think_block = processed_answer["in_think_block"]
                think_buffer = processed_answer["think_buffer"]
                final_answer = processed_answer["final_answer"]
                
                # 如果有非思考内容，立即返回
                if processed_answer["output"]:
                    yield processed_answer["output"]
        
        else:
            # 处理非流式响应
            try:
                result = response.json()
                logger.debug(f"完整响应: {json.dumps(result, ensure_ascii=False)}")
                
                if result.get("code") == 0:
                    answer_data = result.get("data", {})
                    if isinstance(answer_data, dict):
                        answer = answer_data.get("answer", "")
                        if answer:
                            cleaned = self._clean_ragflow_answer(answer, request_type)
                            
                            # 对于非聊天请求，尝试提取JSON
                            if request_type != "chat":
                                try:
                                    # 尝试从响应中提取JSON
                                    json_match = re.search(r'\{.*\}', cleaned, re.DOTALL)
                                    if json_match:
                                        json_data = json.loads(json_match.group())
                                        # 格式化响应以适应 Open-WebUI
                                        formatted = self._format_for_openwebui(json_data, request_type)
                                        logger.debug(f"formatted: {formatted}")
                                        yield formatted
                                    else:
                                        yield cleaned
                                except json.JSONDecodeError:
                                    logger.warning("无法从响应中提取有效的JSON数据")
                                    yield cleaned
                            
                            # 对于普通聊天请求，直接返回回答
                            else:
                                yield cleaned
                    else:
                        logger.error("响应中没有找到答案")
                        yield "错误: 响应中没有找到答案"
                else:
                    error_msg = result.get("message", "未知错误")
                    logger.error(f"RAGFlow错误: {error_msg}")
                    yield f"错误: {error_msg}"
            except json.JSONDecodeError as e:
                logger.error(f"JSON解析错误: {e}")
                logger.error(f"响应内容: {response.text}")
                yield f"响应解析错误: {e}"

    def _clean_ragflow_answer(self, answer: str, request_type: str = "chat") -> str:
        """清理 RAGFlow 的回答内容，彻底移除所有思考标签"""
        if not isinstance(answer, str):
            return ""
        
        # 移除所有思考标签及其内容（包括不完整的标签）
        # 第一遍：移除完整的<think>...</think>标签
        think_pattern = r'<think>.*?</think>'
        cleaned = re.sub(think_pattern, '', answer, flags=re.DOTALL)
        
        # 第二遍：移除可能的不完整<think>标签及其后所有内容
        think_pattern = r'<think>.*'
        cleaned = re.sub(think_pattern, '', cleaned, flags=re.DOTALL)
        
        # 第三遍：移除任何残留的<think>或</think>标签
        cleaned = cleaned.replace("<think>", "").replace("</think>", "")
        
        # 对于非聊天请求，保留JSON格式
        if request_type != "chat":
            # 移除JSON代码块标记
            cleaned = cleaned.replace("```json", "").replace("```", "").strip()
            
            # 尝试提取JSON部分
            json_match = re.search(r'\{.*\}', cleaned, re.DOTALL)
            if json_match:
                cleaned = json_match.group()
        
        # 移除多余的空白行
        cleaned = re.sub(r'\n\s*\n', '\n\n', cleaned)
        
        # 移除首尾空白
        cleaned = cleaned.strip()
        
        return cleaned

    def _format_for_openwebui(self, data: Dict[str, Any], request_type: str) -> str:
        """格式化响应以适应 Open-WebUI 的期望格式"""
        if request_type == "generate_follow_ups":
            if "follow_ups" in data:
                return json.dumps({"suggestions": data["follow_ups"]})
            elif "suggestions" in data:
                return json.dumps({"suggestions": data["suggestions"]})
            else:
                # 如果没有找到预期的键，尝试使用第一个可用的数组
                for key, value in data.items():
                    if isinstance(value, list) and all(isinstance(item, str) for item in value):
                        return json.dumps({"suggestions": value})
        
        elif request_type == "generate_tags":
            if "tags" in data:
                return json.dumps({"tags": data["tags"]})
            else:
                # 如果没有找到预期的键，尝试使用第一个可用的数组
                for key, value in data.items():
                    if isinstance(value, list) and all(isinstance(item, str) for item in value):
                        return json.dumps({"tags": value})
        
        # 默认情况下，返回原始数据的 JSON 表示
        return json.dumps(data)